package nutcore

import chisel3._
import chisel3.util._
import chisel3.util.experimental._
import chiseltest._
import chiseltest.formal._
import chiseltest.experimental._
import chisel3.stage.{ChiselStage, ChiselGeneratorAnnotation}


import chiseltest.formal._
import chiseltest.experimental._
import chisel3.stage.{ChiselStage, ChiselGeneratorAnnotation}

import org.scalatest.Tag
import org.scalatest.flatspec.AnyFlatSpec

import chiseltest.ChiselScalatestTester
import chiseltest.formal.backends.FormalEngineAnnotation

import utils._
import bus.simplebus._

// Available solver engine annotations:
// 
// Z3EngineAnnotation
// CVC4EngineAnnotation
// BtormcEngineAnnotation
// 
// It's been tested that Z3EngineAnnotation cannot work with VecInit. You should
// use BtormcEngineAnnotation in most of circumstances.

class FormalCacheSpec extends AnyFlatSpec with ChiselScalatestTester with Formal {
  
  "Cache Stage 1" should "pass" in {
    verify(new CacheStage1Spec, Seq(BoundedCheck(100), BtormcEngineAnnotation))
  }

  "Cache Stage 2" should "pass" in {
    verify(new CacheStage2Spec, Seq(BoundedCheck(20), BtormcEngineAnnotation))
  }

  "Simultaneity of Arrival of metaRead request & pipelined Req" should "pass" in {
    verify(new CacheFrontStageSpec, Seq(BoundedCheck(20), BtormcEngineAnnotation))
  }

  "Cache Stage 3 Level 1" should "pass" in {
    verify(new CacheStage3SpecLevel1, Seq(BoundedCheck(20), BtormcEngineAnnotation))
  }

  "Cache Stage 3 Level 2" should "pass" in {
    verify(new CacheStage3SpecLevel2, Seq(BoundedCheck(20), BtormcEngineAnnotation))
  }
  
  "Cache Level 1" should "pass" in {
    verify(new SingleLevelCacheSpec, Seq(BoundedCheck(40), BtormcEngineAnnotation))
  }
}


// case class CacheConfig (
//   ro: Boolean = false,
//   name: String = "cache",
//   userBits: Int = 0,
//   idBits: Int = 0,
//   cacheLevel: Int = 1,

//   totalSize: Int = 32, // Kbytes
//   ways: Int = 4
// )

class CacheStage1Spec() extends Module {

  implicit val cacheConfig = new CacheConfig()

  val port = IO(new CacheStage1IO())

  class CacheStage1Wrap(implicit val cacheConfig: CacheConfig) extends Module with HasNutCoreConst with HasCacheConst{
    
    val cash = Module(new CacheStage1())
    val port = IO(new CacheStage1IO())
    cash.io <> port

    // read meta array and data array
    val addr = port.in.bits.addr.asTypeOf(addrBundle)
    val readBusValid = port.in.valid && port.out.ready

    assert(port.metaReadBus.req.bits.setIdx === addr.index)
    assert(port.dataReadBus.req.bits.setIdx === Cat(addr.index, addr.wordIndex))
    assert(port.metaReadBus.req.valid === readBusValid)
    assert(port.dataReadBus.req.valid === readBusValid)

    assert(port.out.bits.req === port.in.bits)
    assert(port.out.valid === (port.in.valid && port.metaReadBus.req.ready && port.dataReadBus.req.ready))
    assert(port.in.ready === ((!port.in.valid || port.out.fire) && port.metaReadBus.req.ready && port.dataReadBus.req.ready))

  }

  val mod = Module(new CacheStage1Wrap())
  mod.port <> port

}

class CacheStage2Spec() extends Module {

  implicit val cacheConfig = new CacheConfig()
  val port = IO(new CacheStage2IO())


  class CacheStage2Wrap(implicit val cacheConfig: CacheConfig) extends Module with HasNutCoreConst with HasCacheConst{        

      val cash = Module(new CacheStage2())
      val port = IO(new CacheStage2IO())
      cash.io <> port

      val req = port.in.bits.req
      val addr = req.addr.asTypeOf(addrBundle)

      val metaWay = observe(cash.metaWay)
      val metaRead = port.metaReadResp
      val metaSaveReq = port.metaWriteBus.req
      val metaReadHit = VecInit(metaRead.map(m => m.valid && (m.tag === addr.tag))).asUInt
      // With the assumption of hit, we can only infer that the vec count is greater than 0,
      // but not exact 1. Tighter assumptions are needed.
      
      // val forwardWaymask = observe(cash.forwardMeta.waymask.getOrElse("b1".U)).asBools
      // val forwardMetaData = observe(cash.forwardMeta.data)
      
      when (port.in.valid) {
        when(!observe(cash.pickForwardMeta)) {

          // hit is calculated from fresh-read meta
          assert(metaWay === port.metaReadResp)
          
          // Implies a hit from fresh-read meta
          when(PopCount(metaReadHit) > 0.U) {
            assert(port.out.bits.hit);
            assert(port.out.bits.waymask === metaReadHit)
          }
        }.otherwise{

        }
      }
  }

  val mod = Module(new CacheStage2Wrap());
  mod.port <> port
}

class CacheFrontStageSpec extends Module {

  implicit val cacheConfig = new CacheConfig()

  val in = IO(Flipped(new SimpleBusUC))
  val out = IO(Decoupled(new Stage2IO))

  class CacheFrontStageWrap(implicit val cacheConfig: CacheConfig) extends CacheModule{

    val io = IO(new Bundle{
      val in = Flipped(new SimpleBusUC(userBits = userBits, idBits = idBits))
      val flush = Input(UInt(2.W))
      val out = Decoupled(new Stage2IO)
    })

    val s1 = Module(new CacheStage1)
    val s2 = Module(new CacheStage2)

    // No forwarded reading from stage 3
    val metaArray = Module(new SRAMTemplateWithArbiter(nRead = 1, new MetaBundle, set = Sets, way = Ways))
    val dataArray = Module(new SRAMTemplateWithArbiter(nRead = 1, new DataBundle, set = Sets * LineBeats, way = Ways))

    PipelineConnect(s1.io.out, s2.io.in, s2.io.out.fire, io.flush(0))
    
    metaArray.io.r(0) <> s1.io.metaReadBus
    dataArray.io.r(0) <> s1.io.dataReadBus

    metaArray.io.w <> DontCare
    dataArray.io.w <> DontCare

    s2.io.metaReadResp := s1.io.metaReadBus.resp.data
    s2.io.dataReadResp := s1.io.dataReadBus.resp.data
    s2.io.dataWriteBus := DontCare
    s2.io.metaWriteBus := DontCare

    s1.io.in <> io.in.req
    io.in.resp <> DontCare
    io.out <> s2.io.out

    // we are investigating the read request, thus forbid write request for this moment
    // since SRAMTemplateWithArbiter defines SRAMTemplate as single-ported
    assume(!metaArray.io.w.req.valid && !dataArray.io.w.req.valid)
    assume(!observe(metaArray.ram.resetState))

    assume(io.out.ready === true.B && stable(io.out.ready))

    when(past(io.in.req.valid)) {
      assert(past(s1.io.in.bits) === s2.io.in.bits.req)
      assert(past(s1.io.in.bits) === s2.io.in.bits.req)

      assert(s2.io.in.valid)  // according to definition of pipelineConnect
      assert(s2.io.out.fire)
      assert(s2.io.in.fire)

      assert(past(s1.io.metaReadBus.req.valid))
      assert(past(s1.io.dataReadBus.req.valid))
      assert(s2.io.in.bits.req.addr === past(s1.io.in.bits.addr))
      assert(getMetaIdx(s2.io.in.bits.req.addr) === s1.io.metaReadBus.resp.addr)
    }
    
    when(io.in.req.valid) {
      assert(s1.io.in.valid)
    }
  }

  val mod = Module(new CacheFrontStageWrap())

  mod.io.in <> in
  mod.io.in.resp <> DontCare 
  mod.io.flush := "b00".U
  out <> mod.io.out
}


// Supports mask in Seq of Bool only. You want to use `.asBools`
// to convert an UInt into Seq of Bool
object MaskOneHotCheck {
  def apply[T <: Data](sel: Seq[Bool], in: Seq[T], out: T, func: T =>T): Unit = {
  }
}

class CacheStage3SpecLevel1 extends Module {

  implicit val cacheConfig = new CacheConfig()
  val port = IO(new CacheStage3IO())

  class CacheStage3Wrap(implicit val cacheConfig: CacheConfig) extends Module with HasNutCoreConst with HasCacheConst{        
  
    val cash = Module(new CacheStage3())
    val port = IO(new CacheStage3IO())
    cash.io <> port

    val req = port.in.bits.req

    val state = observe(cash.state)
    val stateTrans = observe(cash.state2)

    val mmio  = observe(cash.mmio)
    val hit   = observe(cash.hit)
    val miss  = observe(cash.miss)
    val probe = observe(cash.probe)

    // Here we will be enumerating all valid combinations, including
    // - read or write
    // - hit or miss
    when(hit) {      

      when (!mmio) {

        // ==========================================================================
        // READ HIT
        // if read-hit, return the read cacheline directly.
        // This fits for L1 cache.

        assert(!miss)
        assert(port.in.bits.hit)
        assert(port.out.bits.rdata === observe(cash.dataRead))

        // ==========================================================================
        // WRITE HIT
        // the write-hit behavior can be briefly described as:
        // 1. obtain the data and address to be written with / to from the request.
        // 2. merge the data to be written with the data read from cache line.
        // 3. check if current dirty bit of cache data is true, if so:
        //    1. write the data line into memory (update) and clear the data bit
        //    2. write the new mixed data into the cache line.

        val hitWrite = observe(cash.hitWrite)

        when (hitWrite) {
          assert(port.in.bits.req.isWrite())

          assume(!observe(cash.useForwardData))
          val useForwardData = observe(cash.useForwardData)
          when (!useForwardData) {
            assert(observe(cash.dataRead) === observe(cash.dataReadArray))


            assume(PopCount(port.in.bits.waymask) === 1.U)
            val maskBools = port.in.bits.waymask.asBools
            val dataRead = observe(cash.dataRead)
            for (i <- 0 until maskBools.length) {
              when (maskBools(i)) { 
                assert(observe(cash.dataReadArray) === port.in.bits.datas(i).data, s"($i)th assertion violates")
                assert(dataRead === port.in.bits.datas(i).data, s"($i)th assertion violates")
              }
            }

            when (!ro.B) {
              assert(observe(cash.wordMask) === MaskExpand(req.wmask))
              assert(observe(cash.dataHitWriteBus.req.bits.data).asUInt === MaskData(dataRead.asUInt, req.wdata, MaskExpand(req.wmask)))
            }

          }

          // When writing a byte into the cache line, first read the whole cache line out,
          // then replace with 1 byte with the data to be written. This is done by MaskData
          // 
          // Note that dataHitWrite and metaHitWrite are affected by req.cmd from input.
          // 
          // And hitWrite is end here by writing meta and data into corresponding cachelines.
          val wordMask = observe(cash.wordMask)
          val dataHitWrite = observe(cash.dataHitWriteBus).req.bits.data.asUInt;
          assert(dataHitWrite === ((observe(cash.dataRead.asUInt) & ~wordMask) | (port.in.bits.req.wdata & wordMask)) )
          for (i <- 0 until dataHitWrite.getWidth) {
            assert(dataHitWrite(i) === (observe(cash.dataRead.asUInt)(i) && !wordMask(i)) || 
                                       (port.in.bits.req.wdata(i)        &&  wordMask(i))
            )
          }

          val meta = observe(cash.meta)
          assert(meta === Mux1H(port.in.bits.waymask, port.in.bits.metas))
          when(!observe(cash.meta.dirty)) {
            assert(observe(cash.metaHitWriteBus).req.bits.data === Wire(new MetaBundle).apply(tag = meta.tag, valid = true.B, dirty = (!ro).B))
          }
        }

      }

      // read / write miss always triggers allocation.
      // 
      // The difference between read & write condition:
      // If to access a cache line marked as dirty, then for both read and write request
      // the current dirty line should be written into memory at first. 
      // 
      // assume(cacheLevel == 1)
      assume(!ro.B)

      when (!mmio && !probe) {

        val memReadReqCond  = observe(cash.miss) && !port.flush && !observe(cash.meta).dirty
        val memWriteReqCond = observe(cash.miss) && !port.flush &&  observe(cash.meta).dirty

        val readBeatCount = observe(cash.readBeatCnt.value)
        val writeBeatCount = observe(cash.writeBeatCnt.value)

        // Always in burst-read mode when reading from /writing to memory. This is immediately set when
        // in readMemReq state.
        assert(observe(cash.cmd) === port.mem.req.bits.cmd)
        switch(state) {
          is(1.U) {
            assert(observe(cash.cmd) === SimpleBusCmd.readBurst)
          }
          is (3.U) {
            when(writeBeatCount < (LineBeats - 1).U) {
              assert(observe(cash.cmd) === SimpleBusCmd.writeBurst)
            }
          }
        } 

        // Note on State transition:
        // In prior to 

        // For read miss
        // assert(!(past(state === 0.U && memReadReqCond, 2) && past(port.mem.req.fire)))
        when(miss && past(state === 0.U && memReadReqCond, 2) && past(port.mem.req.fire)) {

          // Note: the miss signal leads the state to s_memReadReq, and switch to s_memReadResp.
          assert(past(state) === 1.U && state === 2.U)
          assert(readBeatCount === past(observe(cash.addr.wordIndex)))
        }

        // Both s_memReadReq and s_memWriteReq eventually modify the read/write beat counter
        // for transmitting data between cacheline and memory.

        when(past(state === 2.U && port.mem.resp.fire) ) {
          assert(observe(cash.afterFirstRead))
          assert(readBeatCount === past(readBeatCount) + 1.U)
        }

        // For write miss      
        val memWriteCond = miss && !port.flush && observe(cash.meta.dirty)
        when (memWriteCond) {
          when (past(state === 3.U && port.mem.req.fire)) {
            when(past(stateTrans === 0.U)) {
              assert(past(port.dataReadBus.req.fire))
              assert(stateTrans === 1.U)
              assert(observe(cash.dataWay) === port.dataReadBus.resp.data)
              assert(past(observe(cash.dataWay)) === port.dataReadBus.resp.data)
            }
            assert(writeBeatCount === past(writeBeatCount) + 1.U)
          }

          when ((state === 3.U) && (writeBeatCount === (LineBeats-1).U)) {
            assert(port.mem.req.bits.cmd === SimpleBusCmd.writeLast)
          }

          when (past(state === 3.U && port.mem.req.fire && port.mem.req.bits.cmd === SimpleBusCmd.writeLast)) {
            assert(state === 4.U)
          }

          when (past(state === 4.U && port.mem.resp.fire)) {
            assert(state === 1.U)
          }

          when (past(observe(cash.readingFirst))) {
            assert(observe(cash.inRdataRegDemand) === past(port.mem.resp.bits.rdata))
          }

          assert(port.out.bits.rdata === observe(cash.inRdataRegDemand))
        }

        // Note:
        // Since the access of memory is totally asynchronous, we may not able to
        // observe the exact value of write-beat counter at specific cycle with 
        // bare definition of the module.
        //  
        // If it is desired to observe the value change of write-beat counter, you
        // either assume the mem request is always on fire (both ready and value)
        // is true (however, it can't be always true due to the definition of 
        // mem.req.valid), or connect this module with an actual memory device which
        // is also compatible with SimpleBus. 
        // 
        // assume(port.mem.req.fire)
        // assert(past(writeBeatCount, beat - 1) === (LineBeats - beat).U)

        // when(past(port.mem.req.fire, 1)) {
        //   assert(past(port.mem.req.bits.isWriteLast(), 1))
        // }

        // Note:
        // The assertion with this form will always PASS, however, due to the nature
        // of asynchronocity, BMC may never cover the whole procedure of transmitting
        // the dataword, and k-induction will give UNKNOWN result. 

        // when (past(state === 0.U && memWriteCond, LineBeats * 3 + 1)) {
        //   for (beat <- 1 to LineBeats - 1) {
        //     when(past(port.mem.req.fire, beat)) {
        //       assert(past(writeBeatCount, beat - 1) === past(writeBeatCount, beat) + 1.U, s"@$beat")
        //     }
        //   }
        // }
      }

    }

    // MMIO

    // Note:
    // mmio behave like write-through, and  should never be cached, henceforece
    // never make a hit. However there is no constraint appear in the original
    // design. Thus this assertion does not stand unless there is further cons-
    // traint from the outside.
    // when(mmio) { assert(!hit) }
    {
      when (past(state === 0.U) && past(!port.flush && !probe && !observe(cash.hitReadBurst) && !miss && mmio)) {
        assert(state === 5.U)
      }

      when (past(state === 5.U && port.mmio.req.fire)) {
        assert(state === 6.U)
      }

      when (past(state === 6.U && port.mmio.resp.fire)) {
        assert(state === 7.U)
      }
    }
  }

  val mod = Module(new CacheStage3Wrap()(CacheConfig(cacheLevel = 1)));
  mod.port <> port
}


class CacheStage3SpecLevel2 extends Module {

  implicit val cacheConfig = new CacheConfig()
  val port = IO(new CacheStage3IO())

  class CacheStage3Wrap(implicit val cacheConfig: CacheConfig) extends Module with HasNutCoreConst with HasCacheConst{        
  
    val cash = Module(new CacheStage3())
    val port = IO(new CacheStage3IO())
    cash.io <> port

    val req = port.in.bits.req

    val state = observe(cash.state)
    val stateTrans = observe(cash.state2)

    val mmio  = observe(cash.mmio)
    val hit   = observe(cash.hit)
    val miss  = observe(cash.miss)
    val probe = observe(cash.probe)

    // Probe
    {
      val readBeatCount = observe(cash.readBeatCnt.value)
      val writeBeatCount = observe(cash.writeBeatCnt.value)
      val releaseLast = observe(cash.releaseLast)
      val respToL1Last = observe(cash.respToL1Last)

      val addr = observe(cash.addr)
      when (past(state === 0.U) && past(probe && port.cohResp.fire)) {
        when (past(hit)) {assert(state === 8.U)}
        assert(readBeatCount === past(addr.wordIndex)) 
      }

      when (past(state === 8.U)) {
        when (past(port.cohResp.fire)) {
          assert(readBeatCount === past(readBeatCount) + 1.U)
        }
        when (past(observe(cash.respToL1Fire))) {
          assert(readBeatCount === past(readBeatCount) + 1.U)
        }
        when (past(probe && port.cohResp.fire && observe(cash.releaseLast))) {
          assert(state === 0.U)
        }
        when (observe(cash.respToL1Fire) && observe(cash.respToL1Last)) {
          assert(state === 0.U)
        } 
      }
    }
  }

  val mod = Module(new CacheStage3Wrap()(CacheConfig(cacheLevel = 2)));
  mod.port <> port
}

class SingleLevelCacheSpec extends Module {

  implicit val cacheConfig = new CacheConfig()
  val port = IO(new CacheIO())

  // Simplest case: data-cache, no coh manager, no dynamic exec 
  class CacheWrap(implicit val cacheConfig: CacheConfig) extends Module with HasNutCoreConst with HasCacheConst with HasCacheIO{        

    val s1 = Module(new CacheStage1)
    val s2 = Module(new CacheStage2)
    val s3 = Module(new CacheStage3)
    val metaArray = Module(new SRAMTemplateWithArbiter(nRead = 1, new MetaBundle, set = Sets, way = Ways))
    val dataArray = Module(new SRAMTemplateWithArbiter(nRead = 2, new DataBundle, set = Sets * LineBeats, way = Ways))


    // we are investigating the read request, thus forbid write request for this moment
    // since SRAMTemplateWithArbiter defines SRAMTemplate as single-ported
    // assume(!metaArray.io.w.req.valid && !dataArray.io.w.req.valid)
    // assume(!observe(metaArray.ram.resetState))

    val arb = Module(new Arbiter(new SimpleBusReqBundle(userBits = userBits, idBits = idBits), hasCohInt + 1))
    arb.io.in(hasCohInt + 0) <> io.in.req

    s1.io.in <> arb.io.out

    PipelineConnect(s1.io.out, s2.io.in, s2.io.out.fire, io.flush(0))
    PipelineConnect(s2.io.out, s3.io.in, s3.io.isFinish, io.flush(1))
    io.in.resp <> s3.io.out
    s3.io.flush := io.flush(1)
    io.out.mem <> s3.io.mem
    io.mmio <> s3.io.mmio
    io.empty := !s2.io.in.valid && !s3.io.in.valid

    io.in.resp.valid := Mux(s3.io.out.valid && s3.io.out.bits.isPrefetch(), false.B, s3.io.out.valid || s3.io.dataReadRespToL1)

    io.out.coh.req.ready := true.B
    io.out.coh.resp := DontCare
    io.out.coh.resp.valid := false.B
    s3.io.cohResp.ready := true.B

    metaArray.io.r(0) <> s1.io.metaReadBus
    dataArray.io.r(0) <> s1.io.dataReadBus
    dataArray.io.r(1) <> s3.io.dataReadBus

    metaArray.io.w <> s3.io.metaWriteBus
    dataArray.io.w <> s3.io.dataWriteBus

    s2.io.metaReadResp := s1.io.metaReadBus.resp.data
    s2.io.dataReadResp := s1.io.dataReadBus.resp.data
    s2.io.dataWriteBus := s3.io.dataWriteBus
    s2.io.metaWriteBus := s3.io.metaWriteBus

    assert(io.in.req === s1.io.in)

    val isForwardMeta = observe(s2.isForwardMeta)

    // assume(!observe(metaArray.ram.resetState))
    // assert(!observe(metaArray.ram.wen))
    // assert(!s2.io.out.bits.req.cmd(0))
    // assert(s3.io.in.bits.req.cmd(0) === 0.U)

    // assert()
    // assert(!s1.io.out.valid)
    // // assert(!s2.io.in.ready)
    // assert(!(s1.io.out.valid && s2.io.in.ready))
    // assert(!s2.io.in.valid)
    // assert(!s2.io.out.valid)
    // // assert(!(s3.io.out.ready && (observe(s3.state) === 0.U && !observe(s3.hitReadBurst)) && !observe(s3.miss) && !observe(s3.probe)))
    // // assert(!s3.io.in.ready)
    // assert(!s2.io.out.valid || !s3.io.in.ready)
    // assert(!(s2.io.out.valid && s3.io.in.ready))
    // assert(!s3.io.in.valid)
    // // assert(!s3.io.in.bits.hit)
    // assert(!s3.io.in.valid || !s3.io.in.bits.hit)
    // assert(!observe(s3.hit))
    // assert(!observe(s3.hit) || !observe(s3.req.isWrite()))
    // // assert(observe(s3.meta).dirty)
    // assert(!observe(s3.hitWrite) || observe(s3.meta).dirty)
    // assert(!observe(s3.metaHitWriteBus).req.valid)
    // assert(!s3.io.metaWriteBus.req.valid)
    // assert(!observe(metaArray.ram.resetState))
    // assert(!observe(metaArray.ram.io.w.req).valid)
    // assert(observe(metaArray.ram.wen))
    // assert(observe(metaArray.ram.resetState))
    // assert(observe(metaArray.ram.resetState) || observe(metaArray.ram.wen))
    // // assert(observe(metaArray.ram.io.r).req.ready)
    // assert(!observe(metaArray.ram.io.r.req.ready))
    // assert(!metaArray.io.r(0).req.ready)
    // assert(!s1.io.metaReadBus.req.ready)
    // assert(!s1.io.metaReadBus.req.ready || !s1.io.dataReadBus.req.ready)
    // assert(s1.io.out.valid === s2.io.in.valid)
    assert(s3.io.metaWriteBus.req.valid === s2.io.metaWriteBus.req.valid)
    assert(observe(s3.metaWriteArb.io.out.valid) === s3.io.metaWriteBus.req.valid)
    // assert(!observe(s3.metaWriteArb.io.out.valid))
    // when(observe(s3.metaWriteArb.io.out.valid)) {
    //   assert(observe(s3.metaHitWriteBus.req.valid))
    // }

    // assert(!(s2.io.in.valid && s2.io.metaWriteBus.req.valid))    
    when(s2.io.in.valid && s2.io.metaWriteBus.req.valid) {
      when (s2.io.metaWriteBus.req.bits.setIdx === getMetaIdx(s2.io.in.bits.req.addr)) {
        assert(isForwardMeta)
      }
    } 
  }

  val mod = Module(new CacheWrap()(CacheConfig(ro = true)));
  mod.io <> port
}
